// websocketserver.h
#ifndef WEBSOCKETSERVER_H
#define WEBSOCKETSERVER_H

#include <QObject>
#include <memory>
#include <thread>
#include <mutex>
#include <unordered_set>
#include <boost/asio.hpp>
#include <boost/beast.hpp>
#include <queue>

namespace beast = boost::beast;
namespace http = beast::http;
namespace websocket = beast::websocket;
namespace asio = boost::asio;
using tcp = boost::asio::ip::tcp;

// WebSocket会话类，处理单个客户端连接
class WebSocketSession : public std::enable_shared_from_this<WebSocketSession> {
public:
    WebSocketSession(tcp::socket socket);
    void start(); // 启动会话
    void send(const std::string& message, bool is_binary = false); // 发送消息，支持二进制和文本
    void close(); // 关闭会话
    bool isClosed() const { return closed_; } // 检查会话是否已关闭

    // 消息处理回调函数类型定义
    typedef std::function<void(std::shared_ptr<WebSocketSession>, const std::string&, bool)> MessageHandler;
    void setMessageHandler(MessageHandler handler); // 设置消息处理器

    // 关闭处理回调函数类型定义
    typedef std::function<void(std::shared_ptr<WebSocketSession>)> CloseHandler;
    void setCloseHandler(CloseHandler handler); // 设置关闭处理器

private:
    void doRead(); // 开始读取数据
    void onRead(beast::error_code ec, std::size_t bytes_transferred); // 读取完成回调

    websocket::stream<tcp::socket> ws_; // WebSocket流
    beast::flat_buffer buffer_; // 数据缓冲区
    MessageHandler messageHandler_; // 消息处理回调
    CloseHandler closeHandler_; // 关闭处理回调
    std::atomic<bool> closed_{false}; // 会话关闭状态标志
};

// WebSocket服务器类，管理所有客户端连接
class WebSocketServer : public QObject {
    Q_OBJECT
public:
    explicit WebSocketServer(QObject *parent = nullptr);
    ~WebSocketServer();

    bool start(quint16 port); // 启动服务器
    void stop(); // 停止服务器
    void broadcast(const std::string& message, bool is_binary = false); // 广播消息给所有客户端
    void broadcast(const std::string& message, bool is_binary, std::shared_ptr<WebSocketSession> excludeSession); // 广播消息，排除指定客户端

signals:
    void messageReceived(const QString& message); // 文本消息接收信号
    void clientConnected(); // 客户端连接信号
    void clientDisconnected(); // 客户端断开信号
    void binaryMessageReceived(const QByteArray& data); // 二进制消息接收信号

private:
    void doAccept(); // 开始接受新连接
    void removeSession(std::shared_ptr<WebSocketSession> session); // 移除会话

    asio::io_context ioc_; // IO上下文
    tcp::acceptor acceptor_; // 接受器
    std::thread ioThread_; // IO线程
    std::unordered_set<std::shared_ptr<WebSocketSession>> sessions_; // 会话集合
    std::mutex sessionsMutex_; // 会话集合互斥锁
    std::atomic<bool> running_{false}; // 服务器运行状态标志
};

#endif // WEBSOCKETSERVER_H
